import json
import csv


def build_unsup_data():
    data = []
    data.extend(read_jsonl_data("liji.jsonl"))
    data.extend(read_jsonl_data("mengzi.jsonl"))
    data.extend(read_sishu("sishu.json"))
    data.extend(read_zhuangzi("zhuangzi.json"))
    writer = open("simcse_unsup.txt", "w", encoding="utf-8")
    writer.writelines([line + "\n" for line in data])


def read_jsonl_data(fname):
    unsup_data = []
    for line in open(fname):
        paragraph = json.loads(line)
        for pair in paragraph["pairs"]:
            unsup_data.append(pair["text"])
            unsup_data.append(pair["comment"])
    return unsup_data


def read_sishu(fname):
    unsup_data = []
    paragraphs = json.load(open(fname))
    for paragraph in paragraphs:
        for p in paragraph["pairs"]:
            unsup_data.append(p[0])
            unsup_data.append(p[1])
    return unsup_data


def read_zhuangzi(fname):
    unsup_data = []
    paragraphs = json.load(open(fname))
    for paragraph in paragraphs:
        for comment in paragraph["comments"]:
            unsup_data.extend(list(comment.values()))
    return unsup_data


def build_simcse_sup():
    data = convert_bge_for_simcse("all_bge/all_bge_train.jsonl")
    with open('simcse_sup.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["sent0", "sent1"])
        for row in data:
            writer.writerow(row)


def convert_bge_for_simcse(fname):
    data_pairs = []
    for line in open(fname):
        pair = json.loads(line)
        data_pairs.append((pair["query"], pair["pos"][0]))
    return data_pairs


if __name__ == "__main__":
    # build_unsup_data()
    build_simcse_sup()
